use tokio_modbus::prelude::*;
use tokio_modbus::server::tcp::Server;
use tokio_modbus::server::Service;
use tokio::net::TcpListener;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::future::Future;
use std::pin::Pin;
use std::io;

/*
    从站定义
*/

// 寄存器结构体
#[derive(Clone)]
struct MockService {
    registers: Arc<Mutex<HashMap<u16, u16>>>,
}

// 构造器，初始化寄存器内容
impl MockService {
    fn new() -> Self {
        let mut m = HashMap::new();
        // 初始化寄存器
        m.insert(0, 1234);
        m.insert(1, 5678);
        m.insert(2, 42);
        m.insert(3, 99);
        Self { 
            registers: Arc::new(Mutex::new(m)) 
        }
    }
}

//  相关方法定义
impl Service for MockService {
    type Request = Request<'static>;
    type Response = Response;
    type Error = std::io::Error;  // 改为 io::Error
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + Sync + 'static>>;

    fn call(&self, req: Self::Request) -> Self::Future {
        // 克隆Arc
        let regs = self.registers.clone();
        
        Box::pin(async move {
            match req {
                Request::ReadInputRegisters(addr, cnt) => {
                    let map = regs.lock().unwrap();
                    let mut out = Vec::with_capacity(cnt as usize);

                    for i in 0..cnt {
                        out.push(map.get(&(addr + i)).copied().unwrap_or(0));
                    }
                    Ok(Response::ReadInputRegisters(out))
                }
                Request::ReadHoldingRegisters(addr, cnt) => {
                    let map = regs.lock().unwrap();
                    let mut out = Vec::with_capacity(cnt as usize);
                    for i in 0..cnt {
                        out.push(map.get(&(addr + i)).copied().unwrap_or(0));
                    }
                    Ok(Response::ReadHoldingRegisters(out))
                }
                Request::WriteSingleRegister(addr, value) => {
                    let mut map = regs.lock().unwrap();
                    if map.contains_key(&addr) {
                        map.insert(addr, value);
                        Ok(Response::WriteSingleRegister(addr, value))
                    } else {
                        Err(io::Error::new(io::ErrorKind::InvalidInput, "IllegalDataAddress"))
                    }
                }
                Request::WriteMultipleRegisters(addr, values) => {
                    let mut map = regs.lock().unwrap();
                    for (i, v) in values.iter().enumerate() {
                        let reg_addr = addr + i as u16;
                        if map.contains_key(&reg_addr) {
                            map.insert(reg_addr, *v);
                        } else {
                            return Err(io::Error::new(io::ErrorKind::InvalidInput, "IllegalDataAddress"));
                        }
                    }
                    Ok(Response::WriteMultipleRegisters(addr, values.len() as u16))
                }
                _ => Err(io::Error::new(io::ErrorKind::InvalidInput, "Unsupported function")),
            }
        })
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::init();

    let addr = "127.0.0.1:5502";
    let listener = TcpListener::bind(addr).await?;
    println!("modbus slave listening on {}", addr);

    // 关键点：把 TcpListener 交给 Server::new
    let server = Server::new(listener);

    // 每个新连接时创建一个 Service 实例
    let on_connected = |stream, client_addr| async move {
        println!("New connection from: {}", client_addr);
        Ok::<_, io::Error>(Some((MockService::new(), stream)))
    };

    // 连接处理过程中的错误
    let on_process_error = |err| {
        eprintln!("Connection handling error: {:?}", err);
    };

    // 交给 server 来 accept 和调度
    server.serve(&on_connected, on_process_error).await?;

    Ok(())
}